{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Copyright 2025 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Imagen Product Recontext - Generation at Scale\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-creative-studio/blob/main/experiments/Imagen_Product_Recontext/imagen_product_recontext_at_scale.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%vertex-ai-creative-studio%2Fmain%2Fexperiments%2FImagen_Product_Recontext%imagen_product_recontext_at_scale.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-creative-studio/main/experiments/Imagen_Product_Recontext/imagen_product_recontext_at_scale.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/bigquery/import?url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-creative-studio/main/experiments/Imagen_Product_Recontext/imagen_product_recontext_at_scale.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/bigquery/v1/32px.svg\" alt=\"BigQuery Studio logo\"><br> Open in BigQuery Studio\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-creative-studio/blob/main/experiments/Imagen_Product_Recontext/imagen_product_recontext_at_scale.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://upload.wikimedia.org/wikipedia/commons/9/91/Octicons-mark-github.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "| | |\n",
        "|-|-|\n",
        "|Author(s) | [Layolin Jesudhass](https://github.com/LUJ20), Isidro De Loera"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Bb83RYQlii9v"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "collapsed": true,
        "id": "B5DyrsvEioFw",
        "outputId": "0c7350c4-fdc1-4eda-b536-471c40f1b52c"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: google-cloud-aiplatform in /root/.local/lib/python3.10/site-packages (1.105.0)\n",
            "Requirement already satisfied: google-api-core!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,<3.0.0,>=1.34.1 in /usr/local/lib/python3.10/dist-packages (from google-api-core[grpc]!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,<3.0.0,>=1.34.1->google-cloud-aiplatform) (2.19.2)\n",
            "Requirement already satisfied: google-auth<3.0.0,>=2.14.1 in /usr/local/lib/python3.10/dist-packages (from google-cloud-aiplatform) (2.27.0)\n",
            "Requirement already satisfied: proto-plus<2.0.0,>=1.22.3 in /usr/local/lib/python3.10/dist-packages (from google-cloud-aiplatform) (1.25.0)\n",
            "Requirement already satisfied: protobuf!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<7.0.0,>=3.20.2 in /usr/local/lib/python3.10/dist-packages (from google-cloud-aiplatform) (4.25.5)\n",
            "Requirement already satisfied: packaging>=14.3 in /usr/local/lib/python3.10/dist-packages (from google-cloud-aiplatform) (24.2)\n",
            "Requirement already satisfied: google-cloud-storage<3.0.0,>=1.32.0 in /usr/local/lib/python3.10/dist-packages (from google-cloud-aiplatform) (2.19.0)\n",
            "Requirement already satisfied: google-cloud-bigquery!=3.20.0,<4.0.0,>=1.15.0 in /usr/local/lib/python3.10/dist-packages (from google-cloud-aiplatform) (3.31.0)\n",
            "Requirement already satisfied: google-cloud-resource-manager<3.0.0,>=1.3.3 in /usr/local/lib/python3.10/dist-packages (from google-cloud-aiplatform) (1.14.0)\n",
            "Requirement already satisfied: shapely<3.0.0 in /usr/local/lib/python3.10/dist-packages (from google-cloud-aiplatform) (2.0.6)\n",
            "Requirement already satisfied: google-genai<2.0.0,>=1.0.0 in /usr/local/lib/python3.10/dist-packages (from google-cloud-aiplatform) (1.26.0)\n",
            "Requirement already satisfied: pydantic<3 in /usr/local/lib/python3.10/dist-packages (from google-cloud-aiplatform) (2.10.4)\n",
            "Requirement already satisfied: typing_extensions in /usr/local/lib/python3.10/dist-packages (from google-cloud-aiplatform) (4.12.2)\n",
            "Requirement already satisfied: docstring_parser<1 in /usr/local/lib/python3.10/dist-packages (from google-cloud-aiplatform) (0.16)\n",
            "Requirement already satisfied: googleapis-common-protos<2.0.dev0,>=1.56.2 in /usr/local/lib/python3.10/dist-packages (from google-api-core!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,<3.0.0,>=1.34.1->google-api-core[grpc]!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,<3.0.0,>=1.34.1->google-cloud-aiplatform) (1.66.0)\n",
            "Requirement already satisfied: requests<3.0.0.dev0,>=2.18.0 in /usr/local/lib/python3.10/dist-packages (from google-api-core!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,<3.0.0,>=1.34.1->google-api-core[grpc]!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,<3.0.0,>=1.34.1->google-cloud-aiplatform) (2.32.3)\n",
            "Requirement already satisfied: grpcio<2.0dev,>=1.33.2 in /usr/local/lib/python3.10/dist-packages (from google-api-core[grpc]!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,<3.0.0,>=1.34.1->google-cloud-aiplatform) (1.69.0)\n",
            "Requirement already satisfied: grpcio-status<2.0.dev0,>=1.33.2 in /usr/local/lib/python3.10/dist-packages (from google-api-core[grpc]!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,<3.0.0,>=1.34.1->google-cloud-aiplatform) (1.62.3)\n",
            "Requirement already satisfied: cachetools<6.0,>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from google-auth<3.0.0,>=2.14.1->google-cloud-aiplatform) (5.5.0)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.10/dist-packages (from google-auth<3.0.0,>=2.14.1->google-cloud-aiplatform) (0.4.1)\n",
            "Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.10/dist-packages (from google-auth<3.0.0,>=2.14.1->google-cloud-aiplatform) (4.9)\n",
            "Requirement already satisfied: google-cloud-core<3.0.0,>=2.4.1 in /usr/local/lib/python3.10/dist-packages (from google-cloud-bigquery!=3.20.0,<4.0.0,>=1.15.0->google-cloud-aiplatform) (2.4.1)\n",
            "Requirement already satisfied: google-resumable-media<3.0.0,>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from google-cloud-bigquery!=3.20.0,<4.0.0,>=1.15.0->google-cloud-aiplatform) (2.7.2)\n",
            "Requirement already satisfied: python-dateutil<3.0.0,>=2.8.2 in /usr/local/lib/python3.10/dist-packages (from google-cloud-bigquery!=3.20.0,<4.0.0,>=1.15.0->google-cloud-aiplatform) (2.8.2)\n",
            "Requirement already satisfied: grpc-google-iam-v1<1.0.0dev,>=0.12.4 in /usr/local/lib/python3.10/dist-packages (from google-cloud-resource-manager<3.0.0,>=1.3.3->google-cloud-aiplatform) (0.14.0)\n",
            "Requirement already satisfied: google-crc32c<2.0dev,>=1.0 in /usr/local/lib/python3.10/dist-packages (from google-cloud-storage<3.0.0,>=1.32.0->google-cloud-aiplatform) (1.6.0)\n",
            "Requirement already satisfied: anyio<5.0.0,>=4.8.0 in /root/.local/lib/python3.10/site-packages (from google-genai<2.0.0,>=1.0.0->google-cloud-aiplatform) (4.9.0)\n",
            "Requirement already satisfied: httpx<1.0.0,>=0.28.1 in /usr/local/lib/python3.10/dist-packages (from google-genai<2.0.0,>=1.0.0->google-cloud-aiplatform) (0.28.1)\n",
            "Requirement already satisfied: tenacity<9.0.0,>=8.2.3 in /usr/local/lib/python3.10/dist-packages (from google-genai<2.0.0,>=1.0.0->google-cloud-aiplatform) (8.5.0)\n",
            "Requirement already satisfied: websockets<15.1.0,>=13.0.0 in /usr/local/lib/python3.10/dist-packages (from google-genai<2.0.0,>=1.0.0->google-cloud-aiplatform) (14.1)\n",
            "Requirement already satisfied: annotated-types>=0.6.0 in /usr/local/lib/python3.10/dist-packages (from pydantic<3->google-cloud-aiplatform) (0.7.0)\n",
            "Requirement already satisfied: pydantic-core==2.27.2 in /usr/local/lib/python3.10/dist-packages (from pydantic<3->google-cloud-aiplatform) (2.27.2)\n",
            "Requirement already satisfied: numpy<3,>=1.14 in /usr/local/lib/python3.10/dist-packages (from shapely<3.0.0->google-cloud-aiplatform) (1.26.4)\n",
            "Requirement already satisfied: exceptiongroup>=1.0.2 in /usr/local/lib/python3.10/dist-packages (from anyio<5.0.0,>=4.8.0->google-genai<2.0.0,>=1.0.0->google-cloud-aiplatform) (1.2.2)\n",
            "Requirement already satisfied: idna>=2.8 in /usr/local/lib/python3.10/dist-packages (from anyio<5.0.0,>=4.8.0->google-genai<2.0.0,>=1.0.0->google-cloud-aiplatform) (3.10)\n",
            "Requirement already satisfied: sniffio>=1.1 in /usr/local/lib/python3.10/dist-packages (from anyio<5.0.0,>=4.8.0->google-genai<2.0.0,>=1.0.0->google-cloud-aiplatform) (1.3.1)\n",
            "Requirement already satisfied: certifi in /usr/local/lib/python3.10/dist-packages (from httpx<1.0.0,>=0.28.1->google-genai<2.0.0,>=1.0.0->google-cloud-aiplatform) (2024.12.14)\n",
            "Requirement already satisfied: httpcore==1.* in /usr/local/lib/python3.10/dist-packages (from httpx<1.0.0,>=0.28.1->google-genai<2.0.0,>=1.0.0->google-cloud-aiplatform) (1.0.7)\n",
            "Requirement already satisfied: h11<0.15,>=0.13 in /usr/local/lib/python3.10/dist-packages (from httpcore==1.*->httpx<1.0.0,>=0.28.1->google-genai<2.0.0,>=1.0.0->google-cloud-aiplatform) (0.14.0)\n",
            "Requirement already satisfied: pyasn1<0.7.0,>=0.4.6 in /usr/local/lib/python3.10/dist-packages (from pyasn1-modules>=0.2.1->google-auth<3.0.0,>=2.14.1->google-cloud-aiplatform) (0.6.1)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil<3.0.0,>=2.8.2->google-cloud-bigquery!=3.20.0,<4.0.0,>=1.15.0->google-cloud-aiplatform) (1.17.0)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0.dev0,>=2.18.0->google-api-core!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,<3.0.0,>=1.34.1->google-api-core[grpc]!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,<3.0.0,>=1.34.1->google-cloud-aiplatform) (3.4.1)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0.dev0,>=2.18.0->google-api-core!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,<3.0.0,>=1.34.1->google-api-core[grpc]!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,<3.0.0,>=1.34.1->google-cloud-aiplatform) (2.3.0)\n"
          ]
        }
      ],
      "source": [
        "!pip install --upgrade --user google-cloud-aiplatform\n",
        "\n",
        "from google import genai\n",
        "from google.genai import types\n",
        "from google.cloud import aiplatform, storage\n",
        "from google.cloud.aiplatform.gapic import PredictResponse\n",
        "from google.colab import auth\n",
        "\n",
        "import base64\n",
        "import io\n",
        "import re\n",
        "import timeit\n",
        "import os\n",
        "import json\n",
        "from pathlib import Path\n",
        "from typing import Any, Dict, List, Generator\n",
        "from PIL import Image\n",
        "import matplotlib.pyplot as plt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5BvWL23P18xz",
        "outputId": "e9569d09-f9ef-4211-d26b-ec5b20857103"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "WARNING: google.colab.auth.authenticate_user() is not supported in Colab Enterprise.\n",
            "Authenticated with Google Cloud\n"
          ]
        }
      ],
      "source": [
        "# --------------------------------------------------\n",
        "# Authenticate your Colab session with GCP\n",
        "from google.colab import auth\n",
        "auth.authenticate_user()\n",
        "print(\"Authenticated with Google Cloud\")\n",
        "# --------------------------------------------------\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CTEkagDKitdY"
      },
      "source": [
        "## Helper Functions\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "RgR5sLZbi0FK"
      },
      "outputs": [],
      "source": [
        "#Core Helper Functions for GCS scanning and image display\n",
        "\n",
        "# GCS‑scanning + debug prints\n",
        "\n",
        "def get_mime_type(uri: str) -> str:\n",
        "    ext = os.path.splitext(uri)[1].lower()\n",
        "    if ext == \".png\":\n",
        "        return \"image/png\"\n",
        "    elif ext in (\".jpg\", \".jpeg\"):\n",
        "        return \"image/jpeg\"\n",
        "    else:\n",
        "        raise ValueError(f\"Unsupported extension: {ext}\")\n",
        "\n",
        "def discover_product_batches(input_gcs_uri: str) -> Generator[Dict[str, object], None, None]:\n",
        "    \"\"\"\n",
        "    Yields dicts with:\n",
        "      - product_folder (e.g. \"product_5\")\n",
        "      - image_parts   (List[types.Part])\n",
        "      - product_uris  (List[str])\n",
        "    \"\"\"\n",
        "    m = re.match(r\"gs://([^/]+)/(.+?)/?$\", input_gcs_uri)\n",
        "    if not m:\n",
        "        raise ValueError(f\"Invalid GCS URI: {input_gcs_uri}\")\n",
        "    bucket_name, base_prefix = m.groups()\n",
        "    prefix = base_prefix.rstrip(\"/\") + \"/\"\n",
        "    client = storage.Client()\n",
        "    print(f\"🔍 Scanning bucket={bucket_name} prefix={prefix}\")\n",
        "\n",
        "    # fetch all blobs under prefix\n",
        "    blobs = list(client.list_blobs(bucket_name, prefix=prefix))\n",
        "    print(f\"  • Found {len(blobs)} total objects under {prefix}\")\n",
        "\n",
        "    # group by first‑level folder name\n",
        "    folder_map: Dict[str, List[storage.blob.Blob]] = {}\n",
        "    for b in blobs:\n",
        "        rel = b.name[len(prefix):]  # e.g. \"product_5/thermos_1.png\"\n",
        "        parts = rel.split(\"/\", 1)\n",
        "        if len(parts) != 2:\n",
        "            continue\n",
        "        folder, filename = parts\n",
        "        if not filename.lower().endswith((\".png\", \".jpg\", \".jpeg\")):\n",
        "            continue\n",
        "        folder_map.setdefault(folder, []).append(b)\n",
        "\n",
        "    print(f\"  • Discovered product folders: {list(folder_map.keys())}\")\n",
        "\n",
        "    for folder, blob_list in folder_map.items():\n",
        "        # sort & take up to 3\n",
        "        blob_list.sort(key=lambda b: os.path.basename(b.name))\n",
        "        blob_list = blob_list[:3]\n",
        "\n",
        "        uris = [f\"gs://{bucket_name}/{b.name}\" for b in blob_list]\n",
        "        parts = [\n",
        "            types.Part(\n",
        "                file_data=types.FileData(file_uri=uri, mime_type=get_mime_type(uri))\n",
        "            )\n",
        "            for uri in uris\n",
        "        ]\n",
        "\n",
        "        yield {\n",
        "            \"product_folder\": folder,\n",
        "            \"image_parts\":   parts,\n",
        "            \"product_uris\":  uris,\n",
        "        }\n",
        "\n",
        "# Display helpers\n",
        "def download_gcs_image_bytes(uri: str) -> bytes:\n",
        "    m = re.match(r\"gs://([^/]+)/(.*)\", uri)\n",
        "    if not m:\n",
        "        raise ValueError(f\"Invalid GCS URI: {uri}\")\n",
        "    bucket_name, obj = m.groups()\n",
        "    client = storage.Client()\n",
        "    return client.bucket(bucket_name).blob(obj).download_as_bytes()\n",
        "\n",
        "def prediction_to_pil_image(pred: PredictResponse, size=(640, 640)) -> Image.Image:\n",
        "    b64 = pred[\"bytesBase64Encoded\"]\n",
        "    data = base64.b64decode(b64)\n",
        "    img = Image.open(io.BytesIO(data))\n",
        "    img.thumbnail(size)\n",
        "    return img\n",
        "\n",
        "def display_row(items: List[Any], figsize=(12, 4)):\n",
        "    if not items:\n",
        "        print(\"No items to display.\")\n",
        "        return\n",
        "    fig, axes = plt.subplots(1, len(items), figsize=figsize)\n",
        "    if len(items) == 1:\n",
        "        axes = [axes]\n",
        "    for ax, it in zip(axes, items):\n",
        "        if isinstance(it, Image.Image):\n",
        "            ax.imshow(it)\n",
        "        elif isinstance(it, dict) and \"bytesBase64Encoded\" in it:\n",
        "            ax.imshow(prediction_to_pil_image(it))\n",
        "        else:\n",
        "            ax.text(0.5, 0.5, str(it), ha=\"center\", va=\"center\", wrap=True)\n",
        "        ax.axis(\"off\")\n",
        "    plt.tight_layout()\n",
        "    plt.show()\n",
        "\n",
        "# generate() with fixed SafetySetting keyword args\n",
        "def generate(image_parts: List[types.Part]) -> Dict[str, str]:\n",
        "    import json, re\n",
        "    from google import genai\n",
        "    from google.genai import types\n",
        "\n",
        "    client = genai.Client(\n",
        "        vertexai=True,\n",
        "        project=\"consumer-genai-experiments\",\n",
        "        location=\"global\",\n",
        "    )\n",
        "\n",
        "    user_instr = types.Part.from_text(text=\"\"\"\n",
        "Analyze the provided images of a single product (up to 3). First, identify and describe the product in accurate, natural language: focus on material, color, shape, form, pattern, and distinctive design features.\n",
        "\n",
        "Then, determine a visually appropriate and realistic background or scene where the product would naturally appear and look appealing. Base this on the product’s style and category — for example, place a desk lamp in a home office, or a sneaker in a modern studio.\n",
        "\n",
        "DO NOT GENERATE PEOPLE/CHILDREN\n",
        "Your output should be returned as a JSON object in the format below:\n",
        "{\n",
        "  \"Prompt\": \"<rich description of product and proposed scene>\",\n",
        "  \"product_description\": \"<just the product, no scene>\"\n",
        "}\n",
        "\n",
        "Do not reference the original image’s background or lighting.\n",
        "Do not use placeholders like “in a nice room” — be specific about the setting (e.g., “in a sunlit bohemian-style bedroom with woven textures and indoor plants”).\n",
        "\"\"\")\n",
        "\n",
        "    system_instr = types.Part.from_text(text=\"\"\"\n",
        "Role:\n",
        "You are an expert visual analyst and prompt engineer for AI-based image generation. Your task is to analyze up to 3 input images of the same product, and generate a single, high-quality prompt suitable for AI-driven product image recontextualization.\n",
        "\n",
        "Your objectives are twofold:\n",
        "\n",
        "Describe the product accurately: Identify product category, form, material, texture, color, patterns, and notable design features.\n",
        "\n",
        "Propose a compelling background/scene: Select a suitable environment in which the product would naturally and attractively appear — based on its likely usage, aesthetic, and category.\n",
        "\n",
        "DO NOT GENERATE PEOPLE/CHILDREN\n",
        "\n",
        "Input:\n",
        "Up to 3 images of the same product (e.g., different angles or lighting).\n",
        "No metadata, no background descriptions provided — just images.\n",
        "\n",
        "Output Format:\n",
        "Return a JSON object in this structure:\n",
        "{\n",
        "  \"Prompt\": \"<natural language prompt for recontextualized image generation>\",\n",
        "  \"product_description\": \"<just the product description, no scene>\"\n",
        "}\"\"\")\n",
        "\n",
        "    contents = [ types.Content(role=\"user\", parts=[user_instr, *image_parts]) ]\n",
        "    config = types.GenerateContentConfig(\n",
        "        temperature=0.2,\n",
        "        top_p=0.95,\n",
        "        max_output_tokens=8192,\n",
        "        safety_settings=[\n",
        "            types.SafetySetting(category=\"HARM_CATEGORY_HATE_SPEECH\", threshold=\"OFF\"),\n",
        "            types.SafetySetting(category=\"HARM_CATEGORY_DANGEROUS_CONTENT\", threshold=\"OFF\"),\n",
        "            types.SafetySetting(category=\"HARM_CATEGORY_SEXUALLY_EXPLICIT\", threshold=\"OFF\"),\n",
        "            types.SafetySetting(category=\"HARM_CATEGORY_HARASSMENT\", threshold=\"OFF\"),\n",
        "        ],\n",
        "        system_instruction=[system_instr],\n",
        "        thinking_config=types.ThinkingConfig(thinking_budget=0)\n",
        "    )\n",
        "\n",
        "    output = \"\"\n",
        "    for chunk in client.models.generate_content_stream(\n",
        "        model=\"gemini-2.5-flash\", contents=contents, config=config\n",
        "    ):\n",
        "        output += chunk.text\n",
        "\n",
        "    clean = output.strip()\n",
        "    clean = re.sub(r\"^```(?:\\w+)?\\n\",\"\", clean)\n",
        "    clean = re.sub(r\"\\n```$\",\"\", clean)\n",
        "    return json.loads(clean)\n",
        "\n",
        "\n",
        "# call_product_recontext()\n",
        "def call_product_recontext(\n",
        "    image_bytes_list=None,\n",
        "    image_uris_list=None,\n",
        "    prompt=None,\n",
        "    product_description=None,\n",
        "    disable_prompt_enhancement=True,\n",
        "    sample_count=1,\n",
        "    base_steps=None,\n",
        "    safety_setting=None,\n",
        "    person_generation=None,\n",
        ") -> PredictResponse:\n",
        "    inst: Dict[str, Any] = {\"productImages\": []}\n",
        "    if image_uris_list:\n",
        "        for uri in image_uris_list:\n",
        "            inst[\"productImages\"].append({\"image\": {\"gcsUri\": uri}})\n",
        "    if not inst[\"productImages\"]:\n",
        "        raise ValueError(\"No product images provided.\")\n",
        "    if product_description:\n",
        "        inst[\"productImages\"][0][\"productConfig\"] = {\"productDescription\": product_description}\n",
        "    if prompt:\n",
        "        inst[\"prompt\"] = prompt\n",
        "\n",
        "    params = {\"sampleCount\": sample_count}\n",
        "    if disable_prompt_enhancement: params[\"enhancePrompt\"] = False\n",
        "    if safety_setting:       params[\"safetySetting\"] = safety_setting\n",
        "    if person_generation:    params[\"personGeneration\"] = person_generation\n",
        "    if base_steps:           params[\"baseSteps\"] = base_steps\n",
        "\n",
        "    start = timeit.default_timer()\n",
        "    resp = predict_client.predict(\n",
        "        endpoint=model_endpoint,\n",
        "        instances=[inst],\n",
        "        parameters=params,\n",
        "    )\n",
        "    print(f\"Recontext took {timeit.default_timer()-start:.2f}s\")\n",
        "    return resp\n",
        "\n",
        "# save & upload helper\n",
        "def save_and_upload_recontext_image(\n",
        "    prediction_response,\n",
        "    product_folder: str,\n",
        "    output_bucket_name: str,\n",
        "    output_base_prefix: str = \"cymbal_retail/product_images_output\",\n",
        "    image_index: int = 0,\n",
        "):\n",
        "    out_pref = f\"{output_base_prefix}/{product_folder}\"\n",
        "\n",
        "    if isinstance(prediction_response, list):\n",
        "        prediction_response = prediction_response[0]\n",
        "    img = prediction_to_pil_image(prediction_response)\n",
        "\n",
        "    local = f\"{product_folder}_output_{image_index}.jpg\"\n",
        "    img.convert(\"RGB\").save(local, format=\"JPEG\")\n",
        "    print(f\"Saved local: {local}\")\n",
        "\n",
        "    client = storage.Client()\n",
        "    bucket = client.bucket(output_bucket_name)\n",
        "    # ensure folder exists\n",
        "    if not list(bucket.list_blobs(prefix=out_pref+\"/\")):\n",
        "        bucket.blob(out_pref+\"/\").upload_from_string(\"\", content_type=\"application/x-folder\")\n",
        "        print(f\"🗂 Created folder gs://{output_bucket_name}/{out_pref}/\")\n",
        "\n",
        "    dst = bucket.blob(f\"{out_pref}/{local}\")\n",
        "    dst.upload_from_filename(local, content_type=\"image/jpeg\")\n",
        "    print(f\"Uploaded to gs://{output_bucket_name}/{out_pref}/{local}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VsqJUBnfi6VA"
      },
      "source": [
        "## Initialize Vertex AI Client"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jmYYs65xjIEC",
        "outputId": "b171e916-94ce-4231-c8cd-c6699ac0058f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Prediction client ready\n"
          ]
        }
      ],
      "source": [
        "# Init clients\n",
        "PROJECT_ID = \"consumer-genai-experiments\"\n",
        "LOCATION   = \"us-central1\"\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=LOCATION)\n",
        "predict_client = aiplatform.gapic.PredictionServiceClient(\n",
        "    client_options={\"api_endpoint\": f\"{LOCATION}-aiplatform.googleapis.com\"}\n",
        ")\n",
        "model_endpoint = (\n",
        "    f\"projects/{PROJECT_ID}/locations/{LOCATION}\"\n",
        "    + \"/publishers/google/models/imagen-product-recontext-preview-06-30\"\n",
        ")\n",
        "print(\"Prediction client ready\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qhFfqJ-f7Ofx"
      },
      "source": [
        "# Sequential Run"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "collapsed": true,
        "id": "Sv6j7-gyDP4h",
        "outputId": "ec2634a1-721e-47c8-e7fa-ce232035cd95"
      },
      "outputs": [],
      "source": [
        "# Sequential one by one\n",
        "\n",
        "from datetime import datetime\n",
        "import io\n",
        "from PIL import Image\n",
        "\n",
        "# Define constants\n",
        "INPUT_PREFIX = \"gs://id_test_bucket/cymbal_retail/product_images_input\"\n",
        "OUTPUT_BUCKET = \"id_test_bucket\"\n",
        "\n",
        "# Start time\n",
        "start_time = datetime.now()\n",
        "print(f\"\\nProcess started at: {start_time.strftime('%Y-%m-%d %H:%M:%S')}\")\n",
        "\n",
        "# Discover batches\n",
        "batches = list(discover_product_batches(INPUT_PREFIX))\n",
        "print(f\"\\nTotal product folders discovered: {len(batches)}\")\n",
        "\n",
        "if not batches:\n",
        "    raise RuntimeError(\"No product batches found—check your GCS path & permissions!\")\n",
        "\n",
        "# Process each batch\n",
        "for batch in batches:\n",
        "    print(f\"\\n=== Processing {batch['product_folder']} ===\")\n",
        "\n",
        "    # Preview\n",
        "    imgs = [Image.open(io.BytesIO(download_gcs_image_bytes(u))) for u in batch[\"product_uris\"]]\n",
        "    display_row(imgs)\n",
        "\n",
        "    # Generate prompt & description\n",
        "    gen = generate(batch[\"image_parts\"])\n",
        "    print(\"Prompt:\", gen[\"Prompt\"])\n",
        "    print(\"Desc:  \", gen[\"product_description\"])\n",
        "\n",
        "    # Recontextualize\n",
        "    resp = call_product_recontext(\n",
        "        prompt=gen[\"Prompt\"],\n",
        "        product_description=gen[\"product_description\"],\n",
        "        image_uris_list=batch[\"product_uris\"],\n",
        "        disable_prompt_enhancement=False,\n",
        "        sample_count=1,\n",
        "        safety_setting=\"block_low_and_above\",\n",
        "        person_generation=\"allow_adult\",\n",
        "    )\n",
        "\n",
        "    # Display & upload\n",
        "    preds = list(resp.predictions)\n",
        "    display_row(preds)\n",
        "    for i, p in enumerate(preds):\n",
        "        save_and_upload_recontext_image(\n",
        "            p,\n",
        "            batch[\"product_folder\"],\n",
        "            OUTPUT_BUCKET,\n",
        "            image_index=i\n",
        "        )\n",
        "\n",
        "# End time\n",
        "end_time = datetime.now()\n",
        "duration = end_time - start_time\n",
        "print(f\"\\nProcess completed at: {end_time.strftime('%Y-%m-%d %H:%M:%S')}\")\n",
        "print(f\"⏱Total time taken: {duration}\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yq2D-I0X7KG6"
      },
      "source": [
        "# Parallel threads"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "collapsed": true,
        "id": "CLuQfjXYwpBq",
        "outputId": "9f09bacc-31b8-4258-a452-91ea94320e57"
      },
      "outputs": [],
      "source": [
        "# Parallel threads\n",
        "\n",
        "import concurrent.futures\n",
        "from datetime import datetime  # <-- Added\n",
        "import io\n",
        "from PIL import Image\n",
        "\n",
        "INPUT_PREFIX = \"gs://id_test_bucket/cymbal_retail/product_images_input\"\n",
        "OUTPUT_BUCKET = \"id_test_bucket\"\n",
        "MAX_WORKERS = 5  # Adjust this to control parallelism\n",
        "\n",
        "# Start time\n",
        "start_time = datetime.now()\n",
        "print(f\"\\n🔄 Process started at: {start_time.strftime('%Y-%m-%d %H:%M:%S')}\")\n",
        "\n",
        "batches = list(discover_product_batches(INPUT_PREFIX))\n",
        "print(f\"\\nTotal product folders discovered: {len(batches)}\")\n",
        "\n",
        "if not batches:\n",
        "    raise RuntimeError(\"No product batches found—check your GCS path & permissions!\")\n",
        "\n",
        "def process_batch(batch):\n",
        "    try:\n",
        "        print(f\"\\n=== Processing {batch['product_folder']} ===\")\n",
        "\n",
        "        # Preview\n",
        "        imgs = [Image.open(io.BytesIO(download_gcs_image_bytes(u))) for u in batch[\"product_uris\"]]\n",
        "        display_row(imgs)\n",
        "\n",
        "        # Generate prompt & description\n",
        "        gen = generate(batch[\"image_parts\"])\n",
        "        print(\"Prompt:\", gen[\"Prompt\"])\n",
        "        print(\"Desc:  \", gen[\"product_description\"])\n",
        "\n",
        "        # Recontextualize\n",
        "        resp = call_product_recontext(\n",
        "            prompt=gen[\"Prompt\"],\n",
        "            product_description=gen[\"product_description\"],\n",
        "            image_uris_list=batch[\"product_uris\"],\n",
        "            disable_prompt_enhancement=False,\n",
        "            sample_count=1,\n",
        "            safety_setting=\"block_low_and_above\",\n",
        "            person_generation=\"allow_adult\",\n",
        "        )\n",
        "\n",
        "        # Display & upload\n",
        "        preds = list(resp.predictions)\n",
        "        display_row(preds)\n",
        "        for i, p in enumerate(preds):\n",
        "            save_and_upload_recontext_image(\n",
        "                p,\n",
        "                batch[\"product_folder\"],\n",
        "                OUTPUT_BUCKET,\n",
        "                image_index=i\n",
        "            )\n",
        "    except Exception as e:\n",
        "        print(f\"Error processing batch {batch['product_folder']}: {e}\")\n",
        "\n",
        "# Parallel execution\n",
        "with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:\n",
        "    futures = [executor.submit(process_batch, batch) for batch in batches]\n",
        "    concurrent.futures.wait(futures)\n",
        "\n",
        "# End time\n",
        "end_time = datetime.now()\n",
        "duration = end_time - start_time\n",
        "print(f\"\\nProcess completed at: {end_time.strftime('%Y-%m-%d %H:%M:%S')}\")\n",
        "print(f\"Total time taken: {duration}\")\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "Imagen_Product_Recontext-Generation.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
